import sys
from PyQt5.QtWidgets import (
    QApplication,
    QMainWindow,
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QPushButton,
    QLineEdit,
    QFrame,
    QGridLayout,
    QComboBox,
    QLabel,
    QFileDialog,
    QMessageBox,
    QDialog,
    QRadioButton,
)
from PyQt5.QtWidgets import QPlainTextEdit
from PyQt5.QtGui import QFontMetrics, QPixmap, QImage
from PyQt5.QtCore import Qt
from PyQt5 import QtWidgets, QtGui
from PIL import Image
from wr import *
from wr_lsb import *
from wr_dct import *
import os
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding


class Page2(QMainWindow):
    def __init__(self, parent=None):
        super(Page2, self).__init__(parent)
        self.setWindowTitle("stable diffusion水印系统")
        self.setGeometry(100, 100, 1000, 800)

        # 创建中心部件
        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        # 创建2x2网格布局
        grid_layout = QGridLayout(central_widget)
        grid_layout.setContentsMargins(0, 0, 0, 0)  # 设置边距为0

        # 创建三个模块并添加到网格布局中
        module1 = self.create_module_1()
        module2 = self.create_module_2()
        module3 = self.create_module_3()

        # 设置行和列的拉伸比例
        grid_layout.setRowStretch(0, 1)  # 第一行总宽度拉伸比例为1
        grid_layout.setRowStretch(1, 1)  # 第二行总宽度拉伸比例为1

        # 添加模块到布局中
        grid_layout.addWidget(module1, 0, 0, 1, 1)  # 第一行第一列
        grid_layout.addWidget(module2, 0, 1, 1, 1)  # 第一行第二列
        grid_layout.addWidget(module3, 1, 0, 1, 2)  # 第二行，占满第一列和第二列

    def create_module_1(self):
        # 第一个模块的布局和内容
        image_path = "C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\SD\\LN.jpg"
        widget = QWidget()
        layout = QHBoxLayout(widget)

        # 创建一个QLabel对象用于显示图片
        self.image_label = QLabel()

        # 使用QPixmap加载图片
        pixmap = QtGui.QPixmap(image_path)

        # 如果图片大小大于标签当前大小，则调整标签大小以适应图片
        self.image_label.setPixmap(pixmap.scaled(self.image_label.size(), aspectRatioMode=Qt.KeepAspectRatio))

        # 将QLabel添加到布局中
        layout.addWidget(self.image_label)

        # 设置布局到widget
        widget.setLayout(layout)
        return widget

    def create_module_2(self):
        # 第二个模块的布局和内容
        widget = QWidget()
        layout = QVBoxLayout(widget)

        # 第一行
        layout1 = QHBoxLayout()
        self.pic_path = QLineEdit()
        self.pic_path.setReadOnly(True)  # 设置文本框为只读
        self.pic_path.setText("请选择图片路径")
        # self.pic_path.setStyleSheet("QLineEdit { border: 2px solid black; }")
        self.pic_path.setStyleSheet("QLineEdit { border: 2px solid black; font-size: 12pt; }")
        layout1.addWidget(self.pic_path)
        button = QPushButton("Open File", self)
        button.clicked.connect(self.onButtonClick)
        layout1.addWidget(button)
        layout.addLayout(layout1)

        # 第二行
        # 创建下拉框，并添加选项
        lay1 = QHBoxLayout()
        label = QLabel("差异范围：")
        lay1.addWidget(label)
        self.combo_box = QComboBox()
        self.combo_box.addItems(["0-7", "8-15", "16-31", "32-63", "64-127", "128-255"])
        font = self.combo_box.font()
        font.setPointSize(15)  # 设置字体大小为15点
        self.combo_box.setFont(font)
        lay1.addWidget(self.combo_box)  # 将下拉框添加到布局中
        layout.addLayout(lay1)

        # 第三行
        layout2 = QHBoxLayout()
        label = QLabel("待提取数据长度：")
        layout2.addWidget(label)
        self.data_len = QLineEdit()
        self.data_len.setStyleSheet("QLineEdit { border: 2px solid black; }")
        layout2.addWidget(self.data_len)
        layout.addLayout(layout2)

        # 第三行
        button = QPushButton("点击解密")
        button.clicked.connect(self.on_button_clicked1)
        layout.addWidget(button)

        widget.setLayout(layout)
        return widget

    def create_module_3(self):
        # 第三个模块的布局和内容
        widget = QWidget()
        layout = QVBoxLayout(widget)

        # 创建第一个文本输入框并设置占位提示文字
        self.text_edit1 = QPlainTextEdit()
        self.text_edit1.setPlaceholderText("秘密数据")
        self.text_edit1.setStyleSheet("QPlainTextEdit { border: 2px solid black; }")
        self.text_edit1.setReadOnly(True)

        font_metrics = QFontMetrics(self.text_edit1.document().defaultFont())
        line_height = font_metrics.lineSpacing()

        visible_lines = 16  # 设置可见的行数
        desired_height = line_height * visible_lines
        self.text_edit1.setFixedHeight(desired_height)

        layout.addWidget(self.text_edit1)

        widget.setLayout(layout)
        return widget

    def onButtonClick(self):
        options = QFileDialog.Options()
        fileName, _ = QFileDialog.getOpenFileName(self, "Select a file", "", "All Files (*);;Text Files (*.txt)")
        if not fileName:
            return
        self.pic_path.setText(fileName)
        pil_image = Image.open(fileName)

        try:
            # 将Pillow图像转换为QImage
            qimage = QImage(
                pil_image.tobytes(),
                pil_image.size[0],
                pil_image.size[1],
                QImage.Format.Format_RGB888,
            )
        except Exception as e:
            print("An error occurred:", e)
        # 将QImage转换为QPixmap
        try:
            pixmap = QPixmap.fromImage(qimage)
        except Exception as e:
            print("An error occurred:", e)
        # 在QLabel中显示灰度图片
        self.image_label.setPixmap(pixmap)
        # 在QLabel中显示灰度图片
        self.image_label.setPixmap(pixmap)

    def on_button_clicked1(self):
        # 创建选择对话框
        dialog = QDialog(self)
        dialog.setWindowTitle("选择解密方式")
        dialog.setWindowFlags(dialog.windowFlags() & ~Qt.WindowContextHelpButtonHint)
        dialog.setFixedSize(450, 300)  # 设置对话框的固定大小
        layout = QVBoxLayout(dialog)

        # 添加标签
        # label = QLabel("解密方式选择")
        # label.setAlignment(Qt.AlignCenter)  # 设置标签居中对齐
        # layout.addWidget(label)

        # 创建LSB和PVD的单选按钮
        self.lsb_radio = QRadioButton("LSB")
        self.pvd_radio = QRadioButton("PVD")
        self.dct_radio = QRadioButton("DCT")
        self.lsb_radio.setChecked(True)  # 默认选择LSB

        layout.addWidget(self.lsb_radio)
        layout.addWidget(self.pvd_radio)
        layout.addWidget(self.dct_radio)

        # 添加伸缩项以使按钮居中
        layout.addStretch()
        # 创建确认按钮
        confirm_button = QPushButton("确认")
        confirm_button.clicked.connect(dialog.accept)
        layout.addWidget(confirm_button)

        # 显示对话框并等待用户选择
        if dialog.exec_() == QDialog.Accepted:
            # 获取用户选择的解密方式
            if self.lsb_radio.isChecked():
                extraction_method = "LSB"
            elif self.pvd_radio.isChecked():
                extraction_method = "PVD"
            elif self.dct_radio.isChecked():
                extraction_method = "DCT"
            data = self.data_len.text()
            path = self.pic_path.text()
            if len(data) == 0:
                QMessageBox.warning(self, "警告", "提取数据长度不可为空")
                return
            if not data.isdigit():
                QMessageBox.warning(self, "警告", "提取数据长度必须为整数")
                return
            if not os.path.exists(path):
                QMessageBox.warning(self, "警告", "请先选择图片")
                return
            data = int(data)
            choose = self.combo_box.currentText()
            choose = choose.split("-")
            min_d = int(choose[0])
            max_d = int(choose[1])
            fw = (min_d, max_d)

            # 根据选择的解密方式调用不同的处理函数
            if extraction_method == "LSB":
                s = reader_lsb(path, data)  # 假设您有一个LSB的处理函数
            elif extraction_method == "DCT":
                print(path)
                print(data)
                s = reader_dct(path, data)
            else:
                s = reader(path, data, fw)  # 假设您有一个PVD的处理函数


            # RSA修改

            with open("private_key.pem", "rb") as private_key_file:
                private_key = serialization.load_pem_private_key(private_key_file.read(), password=None)
            # RSA解密
            decrypted_message = private_key.decrypt(
                s,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )



            self.text_edit1.setPlainText("")
            # rsa修改
            self.text_edit1.setPlaceholderText(decrypted_message.decode())


if __name__ == "__main__":
    app = QApplication(sys.argv)
    main_window = Page2()
    main_window.show()
    sys.exit(app.exec_())
